# avprobe
"""
Run avprobe and parse output from it to get information about audio and video files.
"""

import os, subprocess

class Info(object):
    streams = ()
    formats = ()
    
    def __init__(self, filename):
        """
        Run avprobe on a multimedia file and return a report.
        """
        try:
            data = subprocess.check_output(["avprobe", "-show_streams", "-show_format", filename])
        except subprocess.CalledProcessError:
            return None
        self._parse_avp(data)
        self.video = self._get_video()
        self.audio = self._get_audio()
        self.framerate = self._get_framerate()
        self.width = self._get_width()
        self.height = self._get_height()
        self.length = self._get_length()
        
        if self.audio.has_key('codec_name'):
            self.audio_codec = self.audio['codec_name']
        else:
            self.audio_codec = None
        
    def __repr__(self):
        if self.video:
            return "Video: %dx%d, %fs, %ffps." % (self.width, self.height, self.length, self.framerate)
        elif self.audio:
            return "Audio."
        else:
            return "Unknown Multimedia (No Audio or Video)."
        
    def _get_video(self):
        for stream in self.streams:
            if stream['codec_type']=='video':
                return stream
        return None
                
    def _get_audio(self):
        for stream in self.streams:
            if stream['codec_type']=='audio':
                return stream
        return None
        
    def _get_n_video(self):
        n = 0
        for stream in self.streams:
            if stream['codec_type']=='video':
                n += 1
        return n
        
    def _get_n_audio(self):
        n = 0
        for stream in self.streams:
            if stream['codec_type']=='audio':
                n += 1
        return n
        
    def _get_framerate(self):
        if self.video.has_key('r_frame_rate'):
            numerator, denominator = (self.video['r_frame_rate'].split('/') + [1.0])[:2]
            framerate = float(numerator) / float(denominator)
            return framerate
        else:
            return None
            
    def _get_width(self):
        if self.video.has_key('width'):
            return int(self.video['width'])
        else:
            return None
        
    def _get_height(self):
        if self.video.has_key('height'):
            return int(self.video['height'])
        else:
            return None
    
    def _get_length(self):
        if self.video.has_key('duration'):
            return float(self.video['duration'])
        else:
            return None
            
    

    def _parse_avp(self, data):
        """
        Parse avprobe stream and format data.
        """
        junk    = {}
        self.streams = []
        self.formats = []
        current = junk
        lines = data.split('\n')
        for line in lines:
            if '[STREAM]' in line:
                self.streams.append({})
                current = self.streams[-1]
                continue
            if '[FORMAT]' in line:
                self.formats.append({})
                current = self.formats[-1]
                continue
            if '[/STREAM]' in line or '[/FORMAT]' in line:
                current = junk
                continue
            parts = line.split('=')
            if len(parts)==2:
                  key, value = [p.strip() for p in parts]
                  current[key] = value


